#!/usr/bin/python

import xlh
import sys, getopt
from re import finditer
from math import log, ceil, floor
from time import time, ctime

MinG = -10

def specFilter4_100(sp):
	retsp = xlh.MSMS()
	range = sp.mzRange()
	n = 1
	while 1:
		start = range[0] + (n -1) * 100
		end = range[0] + n * 100
		retsp.update(sp.spectrumSlice(start, end).mostIntense(4))
		n += 1
		if start > range[1]:
			break
	retsp.setDescription(sp.description())
	retsp.setCharge(sp.charge())
	retsp.setPrecursorMz(sp.precursorMz())
	return retsp

def xSpectraGen(match):
	massladder = set()
	for i in range(1, match.xStart()):
		massladder.add(match.NPeptide().bMass(i))
	for i in range(1,  1 + len(match.NPeptide().sequence()) - match.xStart()):
		massladder.add(match.NPeptide().yMass(i))
	for i in range(1, match.xEnd()):
		massladder.add(match.CPeptide().bMass(i))
	for i in range(1, 1 + len(match.CPeptide().sequence()) - match.xEnd()):
		massladder.add(match.CPeptide().yMass(i))
	return massladder

def mod2string(mod):
	s = ''
	for (pos, mod) in mod.items():
		s = s + mod + ':' + str(pos) + ';'
	return s[:-1]

def finneganScan(descr):
#	p = descr.find('precIntensity')
	p = descr.find('FinneganScanNumber')
	return descr[p:]
			
def guessXonMatch(match, protTuple, spec):
	thisSpectrum = spec.descriptionSearch(match.source().description()).spectra()[0]
	s = 0
	e = 0
	posToTest = []
	guessed = 0
	aaList = xlh.Aminoacids.sections()
	
	if match.NPeptide().hasXMod():
		# we can set xstart as the position of xlinker
		given = match.NPeptide()
		toGuess = match.CPeptide()
		protein = protTuple[1]
		for (pos, mod) in match.NPeptide().modifications().items():
			if xlh.Modifications.getint(mod, 'xlinker') == 1:
				connections = xlh.Modifications.get(mod, 'links')
				s = pos
	else:
		# we can set xend as the position of xlinker
		given = match.CPeptide()
		toGuess = match.NPeptide()
		protein = protTuple[0]
		for (pos, mod) in match.CPeptide().modifications().items():
			if xlh.Modifications.getint(mod, 'xlinker') == 1:
				connections = xlh.Modifications.get(mod, 'links')
				e = pos
		
	for px in connections.split(','):
		if px in aaList:
			# by now we exclude the last aa... at time of writing the only xlinker
			# used prevents cleavage, so it is impossible to have the last aa modified
			# in a cleaved peptide... it can only be if the peptide is cterminal...
			if toGuess.isCTerminal(protein):
				seq = toGuess.sequence()
			else:
				seq = toGuess.sequence()[:-1]
			for pos in finditer(px, seq):
				posToTest.append(pos.start() + 1)
		elif px == 'np' and toGuess.isNTerminal(protein):
			posToTest.append(1)
		elif px == 'cp' and toGuess.isCterminal(protein):
			posToTest.append(toGuess.length())	

	if not posToTest:
		# there are no possible sites on other peptide
		return False
	elif len(posToTest) == 1:
		# there is only one site, we guess it is the right one
		guessed = posToTest[0]
	else:
		# there is more than one site
		# this is tough, because we now have to choose
		# we may decide this according to the highest 
		# number of ion found...
		# in posToTest I have index n... for each one I should check how many
		# bMass(n-1) and yMass(len-n) ions I have
		matchedIons = 0
		for i in posToTest:
			bIons = map(lambda x: toGuess.bMass(x), range(1, i - 1))
			yIons = map(lambda x: toGuess.yMass(x), range(1, toGuess.length() - i))
			# ok, I'm going weird with map and lambda and so on:
			try:
				nb = len(reduce(lambda x, y: x + y, map(lambda x: thisSpectrum.closeTo(x, Tolerance, TolUnit).peakMasses(), bIons)))
			except TypeError:
				nb = 0
				
			try:
				ny = len(reduce(lambda x, y: x + y, map(lambda x: thisSpectrum.closeTo(x, Tolerance, TolUnit).peakMasses(), yIons)))
			except TypeError:
				ny = 0

			if ny + nb < matchedIons:
				break
			else:
				matchedIons = ny + nb
		### BTW, if there are no matches... be strict, discard this!
		if matchedIons == 0:
			return False
		guessed = i

	if not s:
		s = guessed
	else:
		e = guessed

	match.setXStart(s)
	match.setXEnd(e)

	return True
	
def exitOnError(errmsg, exit):
	sys.stderr.write('The following error(s) occurred:\n')
	sys.stderr.write(errmsg + '\n')
	sys.exit(exit)
def usage():
	print sys.argv[0], """-- a tool to validate OMSSA results in a X-link framework
	
Usage:

xvalidate.py [-h] -d DATABASE -i INFILE -f FILE [-t TOL] [-u TOL unit] [-e PPM] [-R Results type] [-l AA]
			 [-g g-score] [-p p-value] [-v] [-o OUTFILE] [-m mask] [-x] [-A] [-c]

 -h:	Print this help and exit
 -d:	Database used for first MS search (not xlink database!) [Default: nr]
 -i:	File containing OMMSA results [only CSV supported]
 -f:	File with spectra [only MGF supported]
 -t:	Tolerance for ion matching [Default 0.5 Da]
 -u:	Tolerance unit [Default Da]
 -e:	ppm error for precursor match [Default 10]
 -R:	OMSSA results format [csv|oms|omx] [not implemented]
 -l:	Filter peptide length [Default: 3 aa]
 -g:	Filter g-scores under this threshold [Default: 0.0]
 -p:	Filter p-values under this threshold [Default: 0.01]
 -v:	Verbose output 
 -o:	Basename for output files
 -m:	bit mask or value: 4 bit 0000 to mask x-link type
            1   1   1   1	
            |   |   |   +-> has multiple atoms
            |   |   +-----> has different sequences in atoms
            |   +---------> has head/tail and tail/head atoms
            +-------------> has different xlinkers
 -x:	Turn on advanced spectrum filter [Default: False]
 -A:	Print only inferred positive matches [Default: False] (Disabled)
 -c:	csv output [Default: False]
	"""
	sys.exit(1)

OnlyPos = True
XSpecFilter = False
DbIn = 'nr'
ResultsFile = ''
ResultsType = 'csv'
SpectraFile = ''
ShortFilter = 3
Tolerance = 0.5
TolUnit = 'Da'
ppmThreshold = 10
pThreshold = 0.01
gThreshold = 0.0
Verbose = False
Basename = 'default'
bitMask = 0
csvOutput = False

spectrumSampling = {}

try:
	optlist, args = getopt.getopt(sys.argv[1:], 'hd:i:f:R:t:u:l:e:p:g:vo:m:xAc')
except getopt.GetoptError:
	usage()

if not optlist:
	usage()

for o, a in optlist:
	if o == '-h':
		usage()
	if o == '-d':
		DbIn = a
	if o == '-i':
		ResultsFile = a
	if o == '-f':
		SpectraFile = a
	if o == '-l':
		ShortFilter = int(a)
	if o == '-t':
		Tolerance = float(a)
	if o == '-u':
		TolUnit = a
		if TolUnit not in ['Da', 'ppm']:
			exitOnError(TolUnit + " is not a valid unit, asshole!", 1)
	if o == '-e':
		ppmThreshold = float(a)
	if o == '-p':
		pThreshold = float(a)
	if o == '-g':
		gThreshold = float(a)
	if o == '-R':
		ResultsType = a
		if ResultsType not in ['csv', 'oms', 'omx']:
			exitOnError(ResultsType + " is not a valid format", 1)
		exitOnError("Results Type specification hasn't been implemented yet", 1)
	if o == '-v':
		Verbose = True
	if o == '-o':
		Basename = a
	if o == '-x':
		XSpecFilter = True
	if o == '-A':
		OnlyPos = False
	if o == '-c':
		csvOutput = True
	if o == '-m':
		mask = int(a)
		if len(a) == 4:
			try:
				bitMask = int(a, 2)
			except ValueError:
				exitOnError('Wrong bitmask, should be 4-bit', 1)
		elif mask < 0 or mask > 15:
			exitOnError('Wrong bitmask, should be 0<x<15', 1)
			bitMask = mask
		else:
			bitMask = mask
		
		
# do some checks
if not DbIn:
	exitOnError('Missing database name', 1)
if not ResultsFile:
	exitOnError('Missing results file', 1)
if not SpectraFile:
	exitOnError('Missing spectra file', 1)
if ppmThreshold < 0:
	exitOnError('Why you do expect a negative ppm?', 1)
if pThreshold < 0:
	exitOnError('Why you do expect a negative p-value', 1)
#### End of option parsing section

# Define some variables
# chech which modifications are xlinker

XlinkerList = filter(lambda x: xlh.Modifications.getint(x, 'xlinker') == 1, xlh.Modifications.sections())
Acclist = {}	# a dictionary accession -> Protein object
XLinkList = []	# a list of XLink atoms
XEventList = []	# the final list of xlinks!

if Verbose:
	# print operational parameters
	paramOut = "[ Operational Parameters ]\n"
	paramOut += "Database: " + DbIn + '\n'
	paramOut += "OMSSA Results File: " + ResultsFile + '\n'
	paramOut += "Spectra File: " + SpectraFile + '\n'
	paramOut += "Ion Tolerance: " + str(Tolerance) + " " + TolUnit + '\n'
	paramOut += "Precursor Tolerance: " + str(ppmThreshold) + " ppm\n"
	paramOut += "Min. Peptide Length: " + str(ShortFilter) + '\n'
	paramOut += "G-Score Threshold: " + str(gThreshold) + '\n'
	paramOut += "P-Value Threshold: " + str(pThreshold) + '\n'
	paramOut += "Mask: " + str(bitMask) + '\n\n'
#	paramOut += "\nAdvanced Filter: " + str(XspecFilter) + '\n'
#	paramOut += "\nPrint: " + ResultsFile + '\n'
	sys.stderr.write(paramOut)
	

sys.stderr.write('Process started on ' + ctime() + '\n')
now = time()

sys.stderr.write('STEP 1: Read results and filter matches\n')

sys.stderr.write('Reading results file... ')
sys.stderr.flush()

raw = xlh.OMSSAResults()
raw.parseOMSSACsv(ResultsFile)

if Verbose:
	ok = 'OK [ ' + str(raw.length()) + ' matches ]\n'
else:
	ok = 'OK\n'

sys.stderr.write(ok)

sys.stderr.write('Reading spectra file... ')
sys.stderr.flush()
spectra = xlh.MSAnalysis()
spectra.parseMgfFile(SpectraFile)
if Verbose:
	ok = 'OK [ ' + str(spectra.length()) + ' spectra ]\n'
else:
	ok = 'OK\n'
sys.stderr.write(ok)

sys.stderr.write('Checking for data consistency... ')
sys.stderr.flush()
filtered = xlh.OMSSAResults()

# we are trying to parse only matches with a single x-mod... possibly not on the last aa
for m in raw.items():
	
	xcount = 0	# counter for x-modifications in results
	
	# defining NPeptide and CPeptide now is only useful to understand
	# if the x-mod is on the last aminoacid of single peptides
	
	NPeptide = map(lambda x: x.split(':'), m.proteinDescription().split())[0][0]
	if m.peptide().find(NPeptide) == 0:
		CPeptide = map(lambda x: x.split(':'), m.proteinDescription().split())[1][0]
	else:
		CPeptide = map(lambda x: x.split(':'), m.proteinDescription().split())[0][0]
		NPeptide = map(lambda x: x.split(':'), m.proteinDescription().split())[1][0]

	if len(NPeptide) < ShortFilter or len(CPeptide) <= ShortFilter:
		continue

	if NPeptide + CPeptide == m.peptide() or CPeptide + NPeptide == m.peptide():
		# x-peptide sequence should be a combination of single peptides
		# this is the first, slowest test...
		for (pos, mod) in m.modifications().items():
			if mod in XlinkerList:
				# first of all check that modification is not on
				# C-terminal of one of single peptides...
				# this check should be made atomic because I'm not sure
				# every xlinker needs this
				
				# pos, the position, is relative to x-peptide, we should map it
				# on single peptides.
				if pos == len(NPeptide) or pos == len(m.peptide()):
					# len(NPeptide) corresponds to the length of left peptide
					# pos is given starting from 1, so that if the modification is on 
					# the C-terminal of Left peptide, pos == len(NPeptide)
					# for the same reason, pos == len(m.peptide()) means that modification
					# is on C-term of CPeptide...
					# in this case we set xcount to a value greater than 2 and get ot
					# this loop
					xcount = 2
					break # breaks this loop of modifications...
				
				# after all count how many xlinkers we have
				# they should not be more than one per peptide
				xcount += 1
		if xcount == 1:
			# we should find just one x-mod for each peptide
			filtered.addMatch(m)
			
			# get accession for proteins... we can do this later and possibly in a more elegant
			# form, but this saves a lot of time !!!
			for p in map(lambda x: x.split(':'), m.proteinDescription().split())[0][1].split(','):
				Acclist[p] = 1
			for p in  map(lambda x: x.split(':'), m.proteinDescription().split())[1][1].split(','):
				Acclist[p] = 1
	
			
if Verbose:
	ok = 'OK [ ' + str(filtered.length()) + ' matches ]\n'
else:
	ok = 'OK\n'
sys.stderr.write(ok)

del raw

if filtered.length() == 0:
	del spectra
	sys.stderr.write('*** No valid matches for this search! ***\n')
	sys.exit(0)


if Verbose:
	sys.stderr.write('Took ' + str(int(time() - now)) + ' s to complete step 1\n')
	now = time()

sys.stderr.write('STEP 2: Validate remaining matches and protein mapping\n')


sys.stderr.write('Reading protein sequences... ')
sys.stderr.flush()
for p in Acclist.keys():
	Acclist[p] = xlh.Protein(xlh.readEntryFromNCBI(DbIn, p))
if Verbose:
	ok = 'OK [ ' + str(len(Acclist)) + ' sequences ]\n'
else:
	ok = 'OK\n'
sys.stderr.write(ok)

sys.stderr.write('Guessing x-link sites... ')
sys.stderr.flush()

for m in filtered.items():

	proteinList = [[],[]]
	for p in map(lambda x: x.split(':'), m.proteinDescription().split())[0][1].split(','):
		proteinList[0].append(Acclist[p])
	for p in map(lambda x: x.split(':'), m.proteinDescription().split())[1][1].split(','):
		proteinList[1].append(Acclist[p])

	# what remains to do:
	# 1- map every x-link event in a protein framework
	# 2- collect up data and do some summary
	
	tempMatch = xlh.XLinkAtom()
	tempMatch.initWithMatch(m)
	if not tempMatch.orderedPeptide():
		# it could be necessary to reverse protein list, to be consistent with "left" and "right"
		# definitions...
		proteinList.reverse()
	
	# we can define here x-link start and end
	# later we can transform them in protein-relative coordinates...
	
	
	
	for t in [(n, c) for n in proteinList[0] for c in proteinList[1]]:
		
		# t is a tuple with a couple of proteins
		# it is not atomic because a peptide can be found more than once in
		# the protein sequence. We should create a new atomic XLinkAtom instance
		# each time a peptide is within its protein
		# this means: check how many times tempMatch.NPeptide() is within t[0]
		# and how many times tempMatch.CPeptide() is within t[1]

		# before appending the new object we need to apply another filter
		# it happens that I have N-N modifications but peptides are not
		# N-terminal to the protein... this means that we should skip
		# these... actually could be better asking Alessio....

###	remove wrong xlinks according to modification type and its position
###	Update: remove these can be wrong, especially if xlink chemical modification
###			is the same on terminal aa and other aa... Just comment it by now
		# NPeptide:
		invalid = False
		for mod in tempMatch.NPeptide().modifications().values():
			type = xlh.Modifications.get(mod, 'modtype')
			modString = xlh.Modifications.get(mod, 'residues')
			if type == 'np' and not tempMatch.NPeptide().isNTerminal() and tempMatch.NPeptide().sequence()[0] not in modString:
				invalid = True
			if type == 'cp' and not tempMatch.NPeptide().isCTerminal() and tempMatch.NPeptide().sequence()[-1] not in modString:
				invalid = True
		# CPeptide:
		for mod in tempMatch.CPeptide().modifications().values():
			type = xlh.Modifications.get(mod, 'modtype')
			modString = xlh.Modifications.get(mod, 'residues')
			if type == 'np' and not tempMatch.CPeptide().isNTerminal() and tempMatch.CPeptide().sequence()[0] not in modString:
				invalid = True
			if type == 'cp' and not tempMatch.CPeptide().isCTerminal() and tempMatch.CPeptide().sequence()[-1] not in modString:
				invalid = True

		if invalid:
			continue
### end of remove




		PeptideMappings = [[],[]]

		
		offset = 0

		# the following two loops build PeptideMappings list
		# in PeptideMappings[0] we have all NPeptide positions on NProtein
		# in PeptideMappings[1] we have all CPeptide positions on CProtein
		# we have then to generate all the couples, that with all protein couples
		# give the entire set of single, atomic xlinks
		# one step remains, that is: we need to correctly spot x-link position!!!!
		# we may want to do this before all these loops, because it should be the same
		# for all NewXLinks... we should calculate it for tempMatch and then replicate
		# values in all NewXLinks...
		
		# Here on we try to guess the aminoacid involved in x-link. For each atom one
		# peptide is explained by its modification, the other one not. 
		# To do this we should check xlink capabilities in Xlink definition
		
		if not guessXonMatch(tempMatch, t, spectra):
			continue
		
		# Here starts correct position mapping on each protein within each single
		# atomic x-link
		while True:
			found = t[0].sequence()[offset:].find(tempMatch.NPeptide().sequence())
			if found == -1:
				break
			PeptidePosition = offset + found
			offset = PeptidePosition + len(tempMatch.NPeptide().sequence())
			PeptideMappings[0].append(PeptidePosition)

		offset = 0
		while True:
			found = t[1].sequence()[offset:].find(tempMatch.CPeptide().sequence())
			if found == -1:
				break
			PeptidePosition = offset + found
			offset = PeptidePosition + len(tempMatch.CPeptide().sequence())
			PeptideMappings[1].append(PeptidePosition)
			
		for pt in [(np, cp) for np in PeptideMappings[0] for cp in PeptideMappings[1]]:
			
			NewXLink = xlh.XLinkAtom()
			NewXLink.initWithMatch(m)
			NewXLink.setSourceProteins(NProtein = t[0], CProtein = t[1])
			NewXLink.setXStart(tempMatch.xStart())
			NewXLink.setXEnd(tempMatch.xEnd())
			NewXLink.setNMap(pt[0] + NewXLink.xStart())
			NewXLink.setCMap(pt[1] + NewXLink.xEnd())
			
			
			
			XLinkList.append(NewXLink)
			
			if XSpecFilter:
				sp = specFilter4_100(spectra.descriptionSearch(m.description()).spectra()[0])
			else:
				sp = spectra.descriptionSearch(m.description()).spectra()[0]
			NewXLink.setSpectrum(sp)
			NewXLink.calcPValue(Tolerance, TolUnit)
			

#filter on error		
XLinkList = filter(lambda x: abs((x.source().precursorMass() - x.source().precursorTheoMass()) / x.source().precursorMass() * 1000000) <= ppmThreshold , XLinkList)

#filter on p-value
XLinkList = filter(lambda x: x.pValue() <= pThreshold, XLinkList)

if Verbose:
	ok = 'OK [ ' + str(len(XLinkList)) + ' xlinks ]\n'
else:
	ok = 'OK\n'
sys.stderr.write(ok)

if Verbose:
	sys.stderr.write('Took ' + str(int(time() - now)) + ' s to complete step 2\n')
	now = time()

if len(XLinkList) == 0:
	del spectra
	del XLinkList
	sys.stderr.write('*** No x-links for this search! ***\n')
	sys.exit(0)


sys.stderr.write('STEP 3: Build xlink table\n')

sys.stderr.write('Collapsing xlinks... ')
sys.stderr.flush()
XEventList.append(xlh.XLinkEvent())

for m in XLinkList:

	scan = 0
	for x in XEventList:
		if not x.addXlAtom(m):
			scan += 1

	if scan == len(XEventList):

		XEventList.append(xlh.XLinkEvent())
		XEventList[-1].addXlAtom(m)

	# correct spectrumSampling, this will be used later to correct on 
	# G-score caclulation... The idea is that the more a spectrum is matched
	# the less is valuable...
	
	try:
		spectrumSampling[m.spectrum().description()] += 1
	except KeyError:
		spectrumSampling[m.spectrum().description()] = 1

# do some filtering...
XEventList = filter(lambda x: x.gScore(spectrumSampling) <= gThreshold, XEventList)
XEventList = filter(lambda x: x.rank() & bitMask == bitMask, XEventList)


if Verbose:
	ok = 'OK [ ' + str(len(XEventList)) + ' events ]\n' 
else:
	ok = 'OK\n'
sys.stderr.write(ok)

del XLinkList


####
# trying to get some true positives... we take for each rank the first quartile as threshold
# for rank higher than 9 the threshold should not be less than -10

rankDict = {}
thrDict = {}
for r in dict.fromkeys(map(lambda x: x.rank(), XEventList)):
	rankDict[r] = map(lambda x: x.gScore(), filter(lambda x: x.rank() == r, XEventList))
	rankDict[r].sort()
	fi = len(rankDict[r]) / 4.0
	if fi <=1:
		thrIndex = 0
	elif abs(fi * 2 - int(fi) * 2) > 1:
		# Haven't found how to have a nice integer casting...
		thrIndex = int(ceil(fi)) - 1
	else:
		thrIndex = int(fi) - 1
	
	if r >= 8 and rankDict[r][thrIndex] > -10:
		thrDict[r] = -10.0
	elif r >= 2 and rankDict[r][thrIndex] > -8:
		thrDict[r] = -8.0
	elif r >= 0 and rankDict[r][thrIndex] > -5:
		thrDict[r] = -5.0
	else:
		thrDict[r] = rankDict[r][thrIndex]
	
	sys.stderr.write('Threshold for rank ' + str(r) + ': ' + str(thrDict[r]) + '\n') 
####

sys.stderr.write('*** Warning!!! This is just a raw list ***\n')

XEventList.sort(key=lambda obj:obj.gScore())


if csvOutput:
	print 'L-Protein, R-Protein, Start, End, Rank, G-Score, L-Peptide, modstring, R-Peptide, modstring, p-value, error, m/z, charge, number, description'
	for x in XEventList:
		for m in x.xlAtoms():
			ppmerror = (m.source().precursorMass() - m.source().precursorTheoMass()) / m.source().precursorMass() * 1000000
			print ','.join(map(lambda x: str(x), [x.NProtein().accession(), x.CProtein().accession(), x.start(), x.end(), x.rank(), x.gScore(), 
			m.NPeptide().sequence(), mod2string(m.NPeptide().modifications()), 
			m.CPeptide().sequence(), mod2string(m.CPeptide().modifications()), 
			m.pValue(), ppmerror, m.source().precursorMz(), m.source().charge(), 
			m.source().spectrumNumber() + 1, finneganScan(m.spectrum().description())]))
else:
	nev = 1
	for x in XEventList:
		print '[ Event', nev, ']'
		print 'Start Protein:', x.NProtein().accession()
		print 'End Protein:', x.CProtein().accession()
		print 'Start:', x.start(), 'End:', x.end()
		print 'G-Score:', x.gScore(), 'Rank:', x.rank()
		nxl = 1
		if Verbose:
			print 'Xlink atomic events'
			for m in x.xlAtoms():
				ppmerror = (m.source().precursorMass() - m.source().precursorTheoMass()) / m.source().precursorMass() * 1000000
				print '[', nxl, ']'
				print '\tLeft Peptide:', m.NPeptide().sequence(), '[', mod2string(m.NPeptide().modifications()), ']'
				print '\tRight Peptide:', m.CPeptide().sequence(), '[', mod2string(m.CPeptide().modifications()), ']'
				print '\tp-value:', m.pValue()
				print '\tSpectrum Number:', m.source().spectrumNumber() + 1, '[',finneganScan(m.spectrum().description()),']'
				print '\tCharge:', m.source().charge(), '\tm/z:', m.source().precursorMz(), '\tError (ppm):', ppmerror
				nxl += 1
		nev += 1
		print
if Verbose:
	sys.stderr.write('Took ' + str(int(time() - now)) + ' s to complete step 3\n')

sys.stderr.write('Process ended on ' + ctime() + '\n')
#del XEventList
